/**************************************************************************
 *
 * Copyright (C) 2007 University of California, Los Angeles.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 *
 *************************************************************************/

package edu.ucla.cs.typecast.util;

import java.util.*;

/**
 * A tree map that expires entries automatically based on the given time and timeout values.
 * 
 * @param <K>
 * @param <V>
 */

public class ExpiringTreeMap<K, V> extends TreeMap<K, V> {
	/**
	 * Version ID for serialization.
	 */

	private static final long serialVersionUID = 1L;

	/**
	 * Mapping from keys to their expiry times.
	 */

	protected Map<K, Long> expiryTimes = new HashMap<K, Long>();

	/**
	 * The time to perform the next book-keeping process to remove expired entries.
	 */

	protected long nextUpdateTime;

	protected Date customDate = null;

	protected int defaultTimeout;

	public ExpiringTreeMap(Comparator<? super K> c, int defaultTimeout) {
		super(c);
		this.defaultTimeout = defaultTimeout;
	}

	public ExpiringTreeMap(int defaultTimeout) {
		super();
		this.defaultTimeout = defaultTimeout;
	}

	public ExpiringTreeMap(Map<? extends K, ? extends V> m, int defaultTimeout) {
		super(m);
		this.defaultTimeout = defaultTimeout;
	}

	public ExpiringTreeMap(SortedMap<K, ? extends V> m, int defaultTimeout) {
		super(m);
		this.defaultTimeout = defaultTimeout;
	}

	public void clear() {
		super.clear();
		expiryTimes.clear();
	}

	public boolean containsKey(Object key) {
		update();
		return super.containsKey(key);
	}

	public boolean containsValue(Object value) {
		update();
		return super.containsValue(value);
	}

	public Set<Map.Entry<K, V>> entrySet() {
		update();
		return super.entrySet();
	}

	public K firstKey() {
		update();
		return super.firstKey();
	}

	public V get(Object key) {
		update();
		return super.get(key);
	}

	/**
	 * Get the custom date set, if any.
	 * 
	 * @return
	 */

	public Date getCustomDate() {
		return customDate;
	}

	public SortedMap<K, V> headMap(K toKey) {
		update();
		return super.headMap(toKey);
	}

	public boolean isEmpty() {
		update();
		return super.isEmpty();
	}

	public Set<K> keySet() {
		update();
		return super.keySet();
	}

	public K lastKey() {
		update();
		return super.lastKey();
	}

	protected long now() {
		return ((customDate != null) ? customDate : new Date()).getTime();
	}

	/**
	 * Called when a map entry is about to be expired.
	 * 
	 * @param key
	 */

	public void onExpireEntry(K key, V value) {
		// Do nothing
		// Children of this class may override this method
	}

	public V put(K key, V value) {
		return put(key, value, defaultTimeout);
	}

	/**
	 * Associates the specified value with the specified key in this map. The entry will expire
	 * after the given timeout.
	 * 
	 * @param key
	 * @param value
	 * @param timeout
	 * @return
	 */

	public V put(K key, V value, int timeout) {
		// Store the expiring time first

		expiryTimes.put(key, now() + timeout);

		// "Next update" time might move earlier because of this entry

		nextUpdateTime = Math.min(nextUpdateTime, now() + timeout);

		// Add the main entry

		return super.put(key, value);
	}

	public void putAll(Map<? extends K, ? extends V> m) {
		update();
		super.putAll(m);
	}

	public V remove(Object key) {
		update();
		return super.remove(key);
	}

	/**
	 * Specify a custom date used to expire items.
	 * 
	 * @param customDate
	 */

	public void setCustomDate(Date customDate) {
		this.customDate = customDate;
	}

	/**
	 * Specify a custom date used to expire items.
	 * 
	 * @param time
	 */

	public void setCustomDate(long time) {
		setCustomDate(new Date(time));
	}

	public int size() {
		update();
		return super.size();
	}

	public SortedMap<K, V> subMap(K fromKey, K toKey) {
		update();
		return super.subMap(fromKey, toKey);
	}

	public SortedMap<K, V> tailMap(K fromKey) {
		update();
		return super.tailMap(fromKey);
	}

	protected void update() {
		long now = now();
		if (now < nextUpdateTime)
			return;

		synchronized (expiryTimes) {
			Iterator<Map.Entry<K, V>> i = super.entrySet().iterator();
			while (i.hasNext() && i.hasNext()) {
				// Entry of <key, value>

				Map.Entry<K, V> entry = i.next();

				// If the key has expired, remove the entry from both the "main"
				// map and the "last updated time" map

				if (expiryTimes.get(entry.getKey()) <= now) {
					onExpireEntry(entry.getKey(), entry.getValue());
					i.remove();
					expiryTimes.remove(entry.getKey());
				}
			}
		}
	}

	/**
	 * Update the timeout of an entry.
	 * 
	 * @param key
	 */

	protected void update(K key) {
		update(key, defaultTimeout);
	}

	/**
	 * Update the timeout of an entry.
	 * 
	 * @param key
	 * @param timeout
	 */

	protected void update(K key, int timeout) {
		expiryTimes.put(key, now() + timeout);
	}

	public Collection<V> values() {
		update();
		return values();
	}
}
